# -*- coding: utf-8 -*-
# pylint: disable=line-too-long
"""
Implementation of paper [MoA](https://github.com/togethercomputer/MoA).
Here is a simple example for conversation with MoA in Agentscope.
"""
from typing import Optional, Union, Sequence

import agentscope
from agentscope.agents import AgentBase, UserAgent
from agentscope.strategy import MixtureOfAgents
from agentscope.message import Msg


class DialogAgentWithMoA(AgentBase):
    """A simple agent used to perform a dialogue.
    We will show all the modification need to use MoA as the main model"""

    def __init__(
        self,
        name: str,
        moa_module: MixtureOfAgents,  # changed to passing moa_module here
        use_memory: bool = True,
    ) -> None:
        """Initialize the dialog agent.

        Arguments:
            name (`str`):
                The name of the agent.
            sys_prompt (`Optional[str]`):
                The system prompt of the agent, which can be passed by args
                or hard-coded in the agent.
            moa_module (`MixtureOfAgents`):
                The inited MoA module you want to use as the main module.
            use_memory (`bool`, defaults to `True`):
                Whether the agent has memory.
        """
        super().__init__(
            name=name,
            sys_prompt="",
            use_memory=use_memory,
        )
        self.moa_module = moa_module  # change model init to moa_module

    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
        """Reply function of the agent. Processes the input data,
        generates a prompt using the current dialogue memory and system
        prompt, and invokes the language model to produce a response. The
        response is then formatted and added to the dialogue memory.

        Args:
            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
                The input message(s) to the agent, which also can be omitted if
                the agent doesn't need any input.

        Returns:
            `Msg`: The output message generated by the agent.
        """
        # record the input if needed
        if self.memory:
            self.memory.add(x)

        # use the module as below:
        response = self.moa_module(
            Msg("system", self.sys_prompt, role="system"),
            self.memory
            and self.memory.get_memory()
            or x,  # type: ignore[arg-type]
        )

        msg = Msg(self.name, response, role="assistant")

        # Print/speak the message in this agent's voice
        self.speak(msg)

        # Record the message in memory
        if self.memory:
            self.memory.add(msg)

        return msg


if __name__ == "__main__":
    # fill you api keys, or host local models using vllm or ollama.
    model_configs = [
        {
            "config_name": "qwen-max",
            "model_type": "dashscope_chat",
            "model_name": "qwen-max",
            "api_key": "{your_api_key}",
            "generate_args": {
                "temperature": 0.7,
            },
        },
        {
            "config_name": "gemini-pro",
            "model_type": "gemini_chat",
            "model_name": "gemini-pro",
            "api_key": "{your_api_key}",
        },
        {
            "config_name": "gpt-4",
            "model_type": "openai_chat",
            "model_name": "gpt-4",
            "api_key": "{your_api_key}",
            "client_args": {
                "max_retries": 3,
            },
            "generate_args": {
                "temperature": 0.7,
            },
        },
    ]

    agentscope.init(model_configs=model_configs, project="Mixture of Agents")

    user_agent = UserAgent()

    your_moa_module = MixtureOfAgents(
        main_model="qwen-max",  # the models you use
        reference_models=["gpt-4", "qwen-max", "gemini-pro"],
        show_internal=False,  # set to True to see the internal of MoA modules
        rounds=1,  # can range from 0 to inf
    )

    # Init two agents
    dialog_agent = DialogAgentWithMoA(
        name="Assistant",
        moa_module=your_moa_module,
        use_memory=True,  # whether to use memory for this agent
    )
    user_agent = UserAgent()

    # start the conversation between user and assistant
    while True:
        q = user_agent(None)
        if q.content == "exit":  # type "exit" to break the loop
            break
        q = dialog_agent(q)
